#include "config.h"
#include "Proxy.h"

using namespace std;

Proxy::Proxy()
{
	m_is_exit = true;

	// for normal server
	m_listen_port = 3128;
	m_listen_socket = -1;

	// for ssl server
	m_ssl_listen_port = 0;
	m_ssl_listen_socket = -1;
	memset( &m_ssl_ca_certificate, 0, sizeof(x509_cert) );
	memset( &m_ssl_server_certificate, 0, sizeof(x509_cert) );
	memset( &m_ssl_server_privatekey, 0, sizeof(rsa_context) );
	
	// reset thread handle
	memset( &m_hListenPThread, 0, sizeof(pthread_t) );
	memset( &m_hProcessPThread, 0, sizeof(pthread_t) );
	
	// init mutex
	pthread_mutex_init( &m_connection_list_pthread_mutex, NULL );

}

Proxy::~Proxy()
{
	Stop();
	
	// free certificate & private key
	x509_free( &m_ssl_ca_certificate );
	x509_free( &m_ssl_server_certificate );
	rsa_free( &m_ssl_server_privatekey );
}

int Proxy::SetNormalServer( const unsigned short listen_port )
{
	if( !m_is_exit )
	{
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "already started." );
		return -1;
	}
	
	// store ports
	m_listen_port = listen_port;
	
	return 0;
}

int Proxy::SetSSLServer( const unsigned short ssl_listen_port,
	const char* ca_crt_filename, const char* srv_crt_filename, const char* srv_key_filename )
{
	DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_DEBUG,
		"ssl_listen_port: %d, ca_crt_filename: %s, srv_crt_filename: %s, srv_key_filename: %s",
		ssl_listen_port, ca_crt_filename, srv_crt_filename, srv_key_filename );
	
	if( !m_is_exit )
	{
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "already started." );
		return -1;
	}
	
	int ret;

	// reset port to 0
	m_ssl_listen_port = 0;

	// 
	DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_NOTICE, "Loading the certificate and key..." );

	// reset ssl listen port
	m_ssl_listen_port = 0;

    /*
     * 0. Initial seeding of the RNG
     */
	memset( &m_ssl_entropy, 0, sizeof(entropy_context) );
	memset( &m_ssl_ctr_drbg, 0, sizeof(ctr_drbg_context) );
	DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_NOTICE, "Initial seeding of the random generator..." );
    entropy_init( &m_ssl_entropy );
    if( ctr_drbg_init(&m_ssl_ctr_drbg, entropy_func, &m_ssl_entropy, (unsigned char *)SSLCONNECTION_SERVER_RESEEDING, strlen(SSLCONNECTION_SERVER_RESEEDING)) != 0 )
    {
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "failed." );
        return -1;
    }
	DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_NOTICE, "OK." );
	
    /*
     * 1. Load the CA certificate
     */
	memset( &m_ssl_ca_certificate, 0, sizeof(x509_cert) );
	ret = x509parse_crtfile( &m_ssl_ca_certificate, ca_crt_filename );
    if( ret != 0 )
    {
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "unable to parse ca certificate." );
        return -2;
    }


    /*
     * 2. Load sever certificate
     */
	memset( &m_ssl_server_certificate, 0, sizeof(x509_cert) );
	ret = x509parse_crtfile( &m_ssl_server_certificate, srv_crt_filename );
    if( ret != 0 )
    {
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "unable to parse server certificate %d.", ret );
        return -3;
    }

    /*
     * 3. Load the server private RSA key
     */
	memset( &m_ssl_server_privatekey, 0, sizeof(rsa_context) );
    rsa_init( &m_ssl_server_privatekey, RSA_PKCS_V15, 0 );
    ret =  x509parse_keyfile( &m_ssl_server_privatekey, srv_key_filename, NULL );
    if( ret != 0 )
    {
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "unable to parse server private key." );
        return -4;
    }

	DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_NOTICE, "OK." );

	// store ports
	m_ssl_listen_port = ssl_listen_port;
	
	return 0;
}

int Proxy::Start()
{
	int ret_code = 0;

	// check if we already start the server
	Stop();
	
	// set to not closed
	m_is_exit = false;
	
	do
	{
		// check ports to start server
		if( m_listen_port )
		{
			if( StartNormalServer() != 0 )
			{
				DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "StartNormalServer failed." );
				ret_code = -1;
				break;
			}
		}
		if( m_ssl_listen_port )
		{
			if( StartSSLServer() != 0 )
			{
				DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "StartSSLServer failed." );
				ret_code = -2;
				break;
			}
		}
		
		// start connection listen thread
		if( pthread_create(&m_hListenPThread, NULL, ListenThread, (void *)this) )
		{
			DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "create ListenThread failed." );
			ret_code = -3;
			break;
		}
		
		// start connection process thread
		if( pthread_create(&m_hProcessPThread, NULL, ProcessThread, (void *)this) )
		{
			DebugOutput( DEBUG_LEVEL_ERROR, "Start ProcessThread failed." );
			ret_code = -4;
			break;
		}
	} while( false );
	
	// if failed, do stop again
	if( ret_code != 0 ) Stop();

	return 0;
}

int Proxy::Stop()
{
	// set exit flag
	m_is_exit = true;
	
	// wait for thread to stop
	void* status;
	pthread_t zero_thread_t;
	memset( &zero_thread_t, 0, sizeof(pthread_t) );
	if( memcmp(&zero_thread_t, &m_hListenPThread, sizeof(pthread_t)) ) pthread_join( m_hListenPThread, &status );
	if( memcmp(&zero_thread_t, &m_hListenPThread, sizeof(pthread_t)) ) pthread_join( m_hProcessPThread, &status );

	// close normal socket
	if( m_listen_socket != -1 )
	{
		close( m_listen_socket );
		m_listen_socket = -1;
	}
	
	// close ssl socket
	if( m_ssl_listen_socket != -1 )
	{
		close( m_ssl_listen_socket );
		m_ssl_listen_socket = -1;
	}
	
	// free connection list
	if( m_connection_list.size() > 0 )
	{
		vector<Connection*>::iterator it;
		for( it = m_connection_list.begin(); it != m_connection_list.end(); it++ ) delete (*it);
		m_connection_list.clear();
	}
	
	return 0;
}

void* Proxy::ListenThread( void* p_arg )
{
	// get the proxy pointer
	Proxy* p_proxy = (Proxy*)p_arg;
	
	// for timeout
	struct timeval timeout;

	/* master file descriptor list */
	fd_set master;
	/* temp file descriptor list for select() */
	fd_set readfds;

	/* clear the master and temp sets */
	FD_ZERO(&master);
	FD_ZERO(&readfds);

	/* add the listener to the master set */
	if( p_proxy->m_listen_socket != -1 ) FD_SET(p_proxy->m_listen_socket, &master);
	if( p_proxy->m_ssl_listen_socket != -1 ) FD_SET(p_proxy->m_ssl_listen_socket, &master);
	
	// get max fd
	int max_fd = 0;
	if( p_proxy->m_listen_socket > max_fd ) max_fd = p_proxy->m_listen_socket;
	if( p_proxy->m_ssl_listen_socket > max_fd ) max_fd = p_proxy->m_ssl_listen_socket;
	max_fd += 1;
	
	// loop to accept connection
	while( p_proxy != NULL && !p_proxy->m_is_exit )
	{
		// do select
		timeout.tv_sec = 5;
		timeout.tv_usec = 0;
		readfds = master;
		if( select(max_fd, &readfds, NULL, NULL, &timeout) <= 0 ) continue;
		
		// check if fd is set
		if( p_proxy->m_listen_socket != -1 && FD_ISSET(p_proxy->m_listen_socket, &readfds) ) p_proxy->AcceptNormalConnection();
		if( p_proxy->m_ssl_listen_socket != -1 && FD_ISSET(p_proxy->m_ssl_listen_socket, &readfds) ) p_proxy->AcceptSSLConnection();
	}
	
	// indicate we have closed
	p_proxy->m_is_exit = true;
	
	pthread_exit( NULL );
}

void* Proxy::ProcessThread( void* p_arg )
{
	// get the proxy pointer
	Proxy* p_proxy = (Proxy*)p_arg;
	
	// for timeout
	struct timeval timeout;
	
	// for select
	int max_fd = 0;
	fd_set readfds;
	
	// iterator to traverse connection list
	vector<Connection*>::iterator it;

	// loop to accept connection
	while( p_proxy != NULL && !p_proxy->m_is_exit )
	{
		// set readfds
		p_proxy->SetReadFds( &readfds, &max_fd );
		if( max_fd <= 0 )
		{
			// sleep 100 ms, then continue
			usleep( 100000 );
			continue;
		}
		
		// do select
		timeout.tv_sec = 0;
		timeout.tv_usec = 10000;
		if( select(max_fd, &readfds, NULL, NULL, &timeout) <= 0 ) continue;
		
		// lock
		THREAD_MUTEX_LOCK( &p_proxy->m_connection_list_pthread_mutex, "m_connection_list_pthread_mutex" );
		
		// process connection data
		for( it = p_proxy->m_connection_list.begin(); it != p_proxy->m_connection_list.end(); )
		{
			if( (*it)->Process(&readfds) != 0 )
			{
				delete (*it);
				it = p_proxy->m_connection_list.erase( it );
			}
			else it++;
		}
		
		// unlock
		THREAD_MUTEX_UNLOCK( &p_proxy->m_connection_list_pthread_mutex, "m_connection_list_pthread_mutex" );
	}
	
	// indicate we have closed
	p_proxy->m_is_exit = true;
	
	pthread_exit( NULL );
}

int Proxy::StartNormalServer()
{
	// create the listen socket
	m_listen_socket = socket( AF_INET, SOCK_STREAM, IPPROTO_TCP );
	if( m_listen_socket == -1 )
	{
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "Create listen socket failed." );
		return -1;
	}

	// set linger
	struct linger oLinger;
	oLinger.l_onoff = 1;
	oLinger.l_linger = 3; // wait 3 seconds for TCP handshake
	if( setsockopt(m_listen_socket, SOL_SOCKET, SO_LINGER, (char*)&oLinger, sizeof(oLinger)) == -1 )
	{
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "Set socket linger mode failed." );
		return -2;
	}

	// set socket to reuse addr
	int flags = 1;
	setsockopt( m_listen_socket, SOL_SOCKET, SO_REUSEADDR, (const char*)&flags, sizeof(flags) );

	// set to nonblocking (the incoming connection will also be non-blocking)
	flags = fcntl( m_listen_socket, F_GETFL, 0 );
	fcntl( m_listen_socket, F_SETFL, flags|O_NONBLOCK );

	// prepare address to bind
	struct sockaddr_in local_addr;
	memset( &local_addr, 0, sizeof(local_addr) );
	local_addr.sin_family = AF_INET;
	local_addr.sin_port = htons( m_listen_port );
	local_addr.sin_addr.s_addr = htonl( INADDR_ANY );

	// bind the listen socket
	if( bind(m_listen_socket, (struct sockaddr*)&local_addr, sizeof(local_addr)) == -1 )
	{
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "Bind listen socket failed." );
		return -3;
	}
	
	// do listen
	listen( m_listen_socket, 5 );

	return 0;
}

int Proxy::StartSSLServer()
{
	 /*
     * 3. Setup the listening TCP socket
     */
	if( net_bind((int*)&m_ssl_listen_socket, NULL, m_ssl_listen_port) != 0 )
	{
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "Create listen socket failed." );
		return -1;
	}

	// set to nonblocking (the incoming connection will also be non-blocking)
	int flags = fcntl( m_ssl_listen_socket, F_GETFL, 0 );
	fcntl( m_ssl_listen_socket, F_SETFL, flags|O_NONBLOCK );
	
	return 0;
}

int Proxy::AcceptNormalConnection()
{
	int ret_code = 0;
	Connection* p_conn = NULL;
	
	// var to store client address
	struct sockaddr_in client_addr;
	int client_addr_len;

	// var to put client scoket
	int client_sock = -1;

	do
	{
		// accept client connection
		client_addr_len = sizeof(client_addr);
		client_sock = accept( m_listen_socket, (struct sockaddr *)&client_addr, (socklen_t *)&client_addr_len );

		// check connection socket is valid
		if( client_sock == -1 )
		{
			DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "accept connection socket failed." );
			ret_code = -1;
			break;
		}

		// successfuly accept
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_NOTICE, "accept connection: %s:%d", inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port) );

		// create the connection entry
		p_conn = new Connection( client_addr );
		
		// do init
		if( p_conn->Init(client_sock) )
		{
			DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "connection init failed." );
			ret_code = -2;
			break;
		}

		// put it in connection list
		THREAD_MUTEX_LOCK( &m_connection_list_pthread_mutex, "m_connection_list_pthread_mutex" );
		m_connection_list.push_back( p_conn );
		THREAD_MUTEX_UNLOCK( &m_connection_list_pthread_mutex, "m_connection_list_pthread_mutex" );
	} while( false );

	// check if failed
	if( ret_code )
	{
		if( p_conn != NULL ) delete p_conn;
		if( client_sock != -1 ) close( client_sock );
	}
	
	return ret_code;
}

int Proxy::AcceptSSLConnection()
{
	int ret_code = 0;
	SSLConnection* p_conn = NULL;
	
	// var to store client address
	struct sockaddr_in client_addr;
	socklen_t client_addr_len;
	
	// var to put client scoket
	int client_sock = -1;
	
	do
	{
		/*
		* 4 Wait until a client connects
		*/
		if( net_accept(m_ssl_listen_socket, &client_sock, NULL) != 0 )
		{
			DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "net_accept failed." );
			ret_code = -1;
			break;
        }

		// get client address
		client_addr_len = sizeof( client_addr );
		getpeername( client_sock, (struct sockaddr*)&client_addr, &client_addr_len );

		// successfuly accept
		DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_NOTICE, "accept client: %s:%d", inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port) );

		// create the connection entry
		p_conn = new SSLConnection( client_addr );
		if( p_conn->InitSSL(&m_ssl_ctr_drbg, &m_ssl_ca_certificate, &m_ssl_server_certificate, &m_ssl_server_privatekey) != 0 )
		{
			DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "InitSSL failed." );
			ret_code = -2;
			break;
		}

		// reset server's random seed
		if( ctr_drbg_reseed(&m_ssl_ctr_drbg, (unsigned char *)SSLCONNECTION_SERVER_RESEEDING, strlen(SSLCONNECTION_SERVER_RESEEDING)) != 0 )
		{
			DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "server ctr_drbg_reseed failed." );
			ret_code = -3;
			break;
		}
		
		// do init
		if( p_conn->Init(client_sock) )
		{
			DEBUG_OUTPUT_WITH_FUNCTION_NAME( DEBUG_LEVEL_ERROR, "connection init failed." );
			ret_code = -4;
			break;
		}

		// put it in connection list
		THREAD_MUTEX_LOCK( &m_connection_list_pthread_mutex, "m_connection_list_pthread_mutex" );
		m_connection_list.push_back( p_conn );
		THREAD_MUTEX_UNLOCK( &m_connection_list_pthread_mutex, "m_connection_list_pthread_mutex" );
		continue;
	} while( false );
	

	// check if failed
	if( ret_code )
	{
		if( p_conn != NULL ) delete p_conn;
		if( client_sock != -1 ) close( client_sock );
	}
	
	return ret_code;
}

int Proxy::SetReadFds( fd_set* p_read_fds, int* p_max_read_fd )
{
	vector<Connection*>::iterator it;
	
	FD_ZERO( p_read_fds );
	*p_max_read_fd = 0;
	THREAD_MUTEX_LOCK( &m_connection_list_pthread_mutex, "m_connection_list_pthread_mutex" );
	for( it = m_connection_list.begin(); it != m_connection_list.end(); it++ )
		(*it)->SetReadFds( p_read_fds, p_max_read_fd );
	THREAD_MUTEX_UNLOCK( &m_connection_list_pthread_mutex, "m_connection_list_pthread_mutex" );
	
	if( p_max_read_fd > 0 ) p_max_read_fd += 1;
	
	return 0;
}
